# Title     : TODO
# Objective : TODO
# Created by: Administrator
# Created on: 2020/4/28

library(data.table)
library(optparse)
library(ropls)
library(magrittr)
library(tidyverse)
library(ggrepel)
library(ggpubr)

isT <- function(data1, data2) {
  var1 <- var(data1)
  if (var1 == 0) {
    return(F)
  }
  d1Test <- shapiro.test(data1)
  d1P <- d1Test$p.value
  if (d1P < 0.05) {
    return(F)
  }
  var2 <- var(data2)
  if (var2 == 0) {
    return(F)
  }
  d2Test <- shapiro.test(data2)
  d2P <- d2Test$p.value
  if (d2P < 0.05) {
    return(F)
  }
  data <- tibble(value = c(data1, data2), group = factor(c(rep(1, length(data1)), rep(2, length(data2)))))
  bTest <- bartlett.test(value ~ group, data)
  bp <- bTest$p.value
  if (bp < 0.05) {
    return(F)
  }
  return(T)
}

mulTest <- function(data, method) {
  p <- 0
  curMethod <- method
  if (curMethod == "auto") {
    sTest <- data %>%
      group_by(ClassNote) %>%
      summarise(p = {
        var <- var(value)
        if (is.na(var) | var == 0) {
          0
        }else shapiro.test(value)$p.value
      })

    bp <- tryCatch({
      bTest <- bartlett.test(value ~ ClassNote, data)
      bTest$p.value
    }, error = function(e) {
      0
    })

    if (all(sTest$p > 0.05) && bp > 0.05) {
      curMethod <- "anova"
    }else {
      curMethod <- "kw"
    }
  }

  if (curMethod == "kw") {
    test <- kruskal.test(value ~ ClassNote, data = data)
    p <- test$p.value
    if (is.na(p)) {
      p <- 1
    }
    curMethod <- "kruskal.test"
  }else if (curMethod == "anova") {
    test <- compare_means(value ~ ClassNote, data = data, method = "anova")
    p <- test$p
    curMethod <- "anova"
  }
  list(p = p, method = curMethod)
}

groupTest <- function(myData, isPaired, method, sampleInfo) {
  classNotes <- unique(myData$ClassNote)
  group1Name <- classNotes[1]
  group2Name <- classNotes[2]
  group1Sample <- myData %>%
    filter(ClassNote == group1Name) %>%
    .$SampleID
  group2Sample <- myData %>%
    filter(ClassNote == group2Name) %>%
    .$SampleID
  if (isPaired) {
    group1Sample <- sampleInfo %>%
      filter(ClassNote == group1Name) %>%
      arrange(Pair_ID) %>%
      .$SampleID
    group2Sample <- sampleInfo %>%
      filter(ClassNote == group2Name) %>%
      arrange(Pair_ID) %>%
      .$SampleID
  }
  data1 <- myData %>%
    filter(SampleID %in% group1Sample) %>%
    arrange(factor(SampleID, levels = group1Sample)) %>%
    .$value
  data2 <- myData %>%
    filter(SampleID %in% group2Sample) %>%
    arrange(factor(SampleID, levels = group2Sample)) %>%
    .$value
  p <- 0
  curMethod <- method
  if (curMethod == "auto") {
    b <- isT(data1, data2)
    if (b) {
      curMethod <- "t"
    }else {
      curMethod <- "u"
    }
  }

  if (curMethod == "u") {
    test <- wilcox.test(data1, data2, paired = isPaired)
    p <- test$p.value
    curMethod <- "wilcox.test"
  }else if (curMethod == "t") {
    p <- tryCatch({
      test <- t.test(data1, data2, paired = isPaired,var.equal = T)
      test$p.value
    }, error = function(e) {
      1
    })

    curMethod <- "t.test"
  }
  if (is.na(p)) {
    p <- 1
  }
  list(p = p, method = curMethod)
}

myTest <- function(isMulti, myData, configData, sampleInfo) {
  rs <- if (isMulti) {
    mulDiffMethod <- configData["diffMethod.mulMethod", "value"]
    test <- mulTest(myData, mulDiffMethod)
    print(test)
    test
  }else {
    diffMethod <- configData["diffMethod.method", "value"]
    isPaired <- configData["isPaired", "value"] == "T"
    groupTest(myData, isPaired, diffMethod, sampleInfo)
  }
  rs
}

myFc <- function(myData, method) {
  classNotes <- unique(myData$ClassNote)
  group1Name <- classNotes[1]
  group2Name <- classNotes[2]
  group1Sample <- myData %>%
    filter(ClassNote == group1Name) %>%
    .$SampleID
  group2Sample <- myData %>%
    filter(ClassNote == group2Name) %>%
    .$SampleID
  data1 <- myData %>%
    filter(SampleID %in% group1Sample) %>%
    arrange(factor(SampleID, levels = group1Sample)) %>%
    .$value
  data2 <- myData %>%
    filter(SampleID %in% group2Sample) %>%
    arrange(factor(SampleID, levels = group2Sample)) %>%
    .$value
  fc <- 0
  curMethod <- method
  if (curMethod == "auto") {
    b <- isT(data1, data2)
    if (b) {
      curMethod <- "mean"
    }else {
      curMethod <- "median"
    }
  }

  if (curMethod == "median") {
    median1 <- median(data1)
    median2 <- median(data2)
    fc <- median2 / median1
  }else if (curMethod == "mean") {
    mean1 <- mean(data1)
    mean2 <- mean(data2)
    fc <- mean2 / mean1
  }
  list(fc = fc, method = curMethod)
}


option_list <- list(
  make_option("--i", default = "AllMet.csv", type = "character", help = "metabolite data file"),
  make_option("--g", default = "SampleInfo.csv", type = "character", help = "sample group file"),
  make_option("--mc", default = "meta_color.txt", type = "character", help = "metabolite color file"),
  make_option("--sc", default = "sample_color.txt", type = "character", help = "sample color file"),
  make_option("--config", default = "config.csv", type = "character", help = "config file")

)
opt <- parse_args(OptionParser(option_list = option_list))

sampleInfo <- read.csv(opt$g, header = T, stringsAsFactors = F) %>%
  mutate(ClassNote = factor(ClassNote, levels = unique(ClassNote)))

configData <- read.csv(opt$config, header = F, stringsAsFactors = F) %>%
  set_colnames(c("arg", "value")) %>%
  column_to_rownames("arg")

configData

fcMethod <- configData["diffMethod.fcMethod", "value"]
isMulti <- length(unique(sampleInfo$ClassNote)) > 2

sampleInfo

parent <- "./"
fileName <- paste0(parent, "Class_Table.csv")
groups <- unique(sampleInfo$ClassNote)
data <- read_csv(fileName) %>%
  rowwise() %>%
  do({
    result <- as_tibble(.)
    eachData <- result %>%
      select(-c("Class")) %>%
      gather("SampleID", "value") %>%
      inner_join(sampleInfo, by = c("SampleID"))
    for (group in groups) {
      sampleIds <- sampleInfo %>%
        filter(ClassNote == group) %>%
        .$SampleID
      groupData <- result[sampleIds] %>%
        unlist()
      result[, paste0(group, ".Mean")] <- mean(groupData)
      result[, paste0(group, ".Median")] <- median(groupData)
      result[, paste0(group, ".SD")] <- sd(groupData)
      iqr11 <- quantile(groupData, 0.25)
      iqr12 <- quantile(groupData, 0.75)
      iqr1 <- paste0("[", iqr11, ",", iqr12, "]")
      result[, paste0(group, ".IQR")] <- iqr1
    }
    print("===")
    print(result)
    rs <- myTest(isMulti, eachData, configData, sampleInfo)
    print(rs)
    result$P <- rs$p
    result$test.method <- rs$method
    if (!isMulti) {
      result$FC <- myFc(eachData, fcMethod)$fc
    }
    result
  }) %>%
  select(-c(sampleInfo$SampleID))

data

if (!isMulti) {
  data <- data %>%
    mutate(FDR = p.adjust(P, method = "fdr")) %>%
    select(-FC, everything()) %>%
    mutate(log2FC = log(FC, 2)) %>%
    select(-"test.method", everything(), "test.method")
}else {
  data <- data %>%
    mutate(FDR = p.adjust(P, method = "fdr")) %>%
    select(-"test.method", everything(), "test.method")
}

print(data)
write_csv(data, "Class_Barplot_by_Group_Test.csv")


